fix(vllm): never transfer null-block-backed chunks; bound retrieves by admitted external tokens - #46
Open
myshytf wants to merge 4 commits into
Conversation
Add fine sub-block projection and carry external chunk geometry from the vLLM adapter into engine-driven SHM and pickle transfers. Preserve coarse identity behavior and fail before slot reservation on invalid geometry. Co-authored-by: Hermes Agent <hermes-agent@nousresearch.com> Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
Avoid applying a whole-dict fused-attention layout detector to mixed rank-4 attention and rank-3 Mamba tensors. Preserve the legacy physical-block check for homogeneous and single-group registrations while delegating multigroup validation to the existing per-group engine-driven context. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com>
The vLLM scheduler queries the connector with the block-aligned part of a local prefix-cache hit and decides afterwards how many external tokens it loads. When the remote hit does not exceed the local sub-block tail it keeps the local tokens, admits zero external tokens, and schedules the forward pass in the same step. The MP connector ignored that decision: its tracker still held the full lookup hit, so update_state_after_alloc entered WAITING_FOR_LOAD and the worker scattered every looked-up chunk into the request's freshly allocated blocks while the forward pass was already reading and writing them. The tracker now records the admitted external token count. A retrieve covers [num_vllm_hit_tokens, admitted end) rounded up to the object boundary and capped at the lookup hit; zero admitted tokens is no retrieve. Lookup locks on chunks past the admitted range are released at update_state_after_alloc since no retrieve will release them. Trackers that never receive an admission keep the previous behavior. Store accounting is unchanged: num_stored_tokens still counts the lookup hit, so chunks LMCache already holds are not re-stored. Validation: tests/v1/test_vllm_mp_connector_metadata.py and tests/v1/test_lmcache_mp_connector_fine_chunks.py, 30 passed, run with the production image interpreter against the deployed source composition.
vLLM reserves block id 0 as the null placeholder block. Hybrid managers put it into a request's block table wherever the request owns no state: the recurrent checkpoint slots before the block a prefix-cache hit or external load resumes from. Its page bytes are shared by every KV cache group of the same layer index and are assumed to stay zero. The MP connector sliced block ids per group without regard to that placeholder. A retrieve for a resumed hybrid request scattered the recurrent-state objects of the leading chunks into block 0, overwriting the null pages of every group, and a store whose leading chunks were null-backed gathered placeholder bytes and persisted them as recurrent cache content. With the packed fp8_ds_mla MLA reader, whose masked candidates load physical slot 0, the corrupted null page turned every later decode step non-finite. Stores now end before the first chunk whose block in any group is the null block: objects are keyed as a contiguous prefix, so nothing past such a chunk can be published. Retrieves substitute a negative skipped id for the null block; the worker drops the leading skipped chunks of each group and loads the remaining suffix, so the attention group still restores the full range while the recurrent group restores only the checkpoint it resumes from. A skipped destination after a kept chunk fails the retrieve closed so the engine recomputes. Validation: tests/v1/test_vllm_mp_connector_metadata.py, tests/v1/test_lmcache_mp_connector_fine_chunks.py, tests/v1/multiprocess/test_dcp_fine_external_chunks.py and tests/v1/multiprocess/test_engine_driven_recurrent_alias.py, 82 passed with the production image interpreter against the deployed source composition. Existing tests that used block id 0 as an ordinary block now use ids from 1.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (1)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Result
Hybrid (recurrent + attention) requests that resume from a prefix-cache hit or
an external load no longer make the MP connector write into, or read from,
vLLM's null placeholder block, and a local sub-block prefix hit that the
scheduler decides to keep no longer triggers a duplicate external retrieve
that races the forward pass.
Behavior
vLLM reserves block id 0 as the placeholder that pads every block table. Its
page bytes are shared by every KV cache group of the same layer index and are
assumed to stay zero. Hybrid managers put the placeholder into a resumed
request's recurrent block table wherever the request owns no checkpoint, so a
request resumed at 4,608 of 12,288 tokens carries
[0, 0, c2, w]for eachrecurrent group.
GetRetrieveMetadatasubstitutesSKIPPED_BLOCK_ID(-1) for the nullblock. The worker drops the leading skipped chunks of each engine group and
scatters the remaining suffix, so the attention group still restores the
whole range while a recurrent group restores only the checkpoint it resumes
from. A skipped destination after a kept chunk fails the retrieve closed and
the engine recomputes.
GetStoreMetadataends the store window before the first chunk whose blockin any group is the null block. Objects are keyed as a contiguous prefix, so
nothing past that chunk can be published; previously the placeholder bytes
were persisted as recurrent cache content.
update_state_after_allocrecords the external token count the scheduleradmitted. A retrieve covers
[num_vllm_hit_tokens, admitted end)rounded upto the object boundary and capped at the lookup hit; zero admitted tokens is
no retrieve, and lookup locks past the admitted range are released there.
Trackers that never receive an admission keep the previous behavior.
Store accounting is unchanged:
num_stored_tokensstill counts the lookup hit,so chunks LMCache already holds are not re-stored.
Evidence
Kimi-K3 (TP8/DCP8, 1,536-token recurrent cadence, 12,288-token attention
block, LMCache chunk 1,536) with the packed
fp8_ds_mlacache:the engine corrupted for all later requests, including a cold 34-token
prompt. Raw KV page dumps taken before and after the resumed prefill step
showed block 0 going from all zeros to 1,284 non-zero records on every rank,
timed with the LMCache
Retrieved 4608 tokensline; the packed MLA readerloads physical slot 0 for its masked candidates, so the corrupted null page
fed NaN into every later decode step. With the standard
fp8cache the samewrites happen and stay unobserved because the dense reader never touches
block 0.
local-hit handling ([3/N][Core][KV Connector] Support reliable partial-tail KV offload for sub-block prompts vllm-project/vllm#49502): the lookup is shown a
block-aligned local hit of 0, the scheduler then keeps the 4,608-token local
tail and admits 0 external tokens, and
update_state_after_allocignored thatcount. The target logged
Finished recving KV transfer ... in unexpected status RUNNINGon every affected request.Validation
tests/v1/test_vllm_mp_connector_metadata.py,tests/v1/test_lmcache_mp_connector_fine_chunks.py,tests/v1/multiprocess/test_dcp_fine_external_chunks.py,tests/v1/multiprocess/test_engine_driven_recurrent_alias.py,tests/v1/multiprocess/test_engine_driven_transfer.py,tests/v1/test_vllm_kv_cache_groups.py: 138 passed with the production imageinterpreter (
/opt/venv/bin/python -m pytest).the vLLM-facing connector treats 0 as the placeholder.
branch mounted is recorded in the candidate directory
(
kimi-k3-production/candidates/k3-fp8-ds-mla-20260904/QUALIFICATION.md).This branch is stacked on #27 (
feat/dcp-fine-external-release).Review requirement
AI assistance was used to implement and test this change. A human maintainer
must review the retrieve suffix projection (
_drop_skipped_chunksand thestart_token_idxoverride of_group_transfer_inputs), the store windowtrimming, and the admitted-range lock release before the PR is marked ready.